"""
@Filename       : bertweet_encoder.py
@Create Time    : 2021/5/13 9:31
@Author         : Rylynn
@Description    : 

"""

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
#
#
class BertweetEncoder(nn.Module):
    def __init__(self):
        super(BertweetEncoder, self).__init__()
        self.bertweet = AutoModel.from_pretrained("vinai/bertweet-base")
        # For transformers v4.x+:
        self.tokenizer = AutoTokenizer.from_pretrained("vinai/bertweet-base", normalization=True)

    def forward(self, input):
        user_tweets_list = []
        for user_tweets in input:
            tweets = [tweet.strip() for tweet in user_tweets.split('\t')[-1].strip().split('||| ')[1:]]
            for tweet in tweets:
                input_ids = tokenizer.encode(tweet)
                user_tweets_list.append(input_ids)
                print(bertweet(torch.tensor(tweets_ids))[1])
        features = self.bertweet(user_tweets_list)  # Models outputs are now tuples
        return features


if __name__ == '__main__':
    # encoder = BertweetEncoder()
    bertweet = AutoModel.from_pretrained("vinai/bertweet-base")
    tokenizer = AutoTokenizer.from_pretrained("vinai/bertweet-base", normalization=True)
    with open('/root/location/dataset/world/user_info.test/user_info.test', 'r', encoding='utf8') as f:
        for line in f.readlines():
            tweets = [tweet.strip() for tweet in line.split('\t')[-1].strip().split('||| ')[1:]]
            tweets_ids = []

            for tweet in tweets:
                print(tweet)
                input_ids = tokenizer.encode(tweet)
                tweets_ids.append(input_ids)
                print(bertweet(torch.tensor(tweets_ids))[1])
            break

